import websocket, _thread
from json import loads
from jwt import encode
from MyTokens import ConstMyTokens
from time import time, sleep

def GetHuyaToken(RoomId, 
                 AppId, 
                 AppSecret):
    '''
    #########################################
    [功能描述] 获取虎牙令牌.
    [参数列表] 1. RoomId    房间号;
               2. AppId     开发者 ID;
               3. AppSecret 开发者密钥;
    [返回类型] str.
    #########################################
    '''

    intUnixTs = int(time())

    dicPayLoad = {
        "iat": intUnixTs, 
        "exp": intUnixTs + 600, 
        "roomId": RoomId, 
        "appId": AppId
    }
    
    return encode(dicPayLoad, AppSecret, algorithm = "HS256").decode()

def OnMessage(WS, Msg):
    dicMsg = loads(Msg)
    print(dicMsg)

def OnError(WS, Err):
    print(Err)

def OnClose(WS):
    print("### closed ###")

def OnOpen(WS):
    def Run(*args):
        WS.send('''{"command": "subscribeNotice", 
                    "data": ["getMessageNotice", 
                             "getSendItemNotice", 
                             "getOpenNobleNotice", 
                             "getOpenGuardianNotice", 
                             "getUserMutedNotice"], 
                    "reqId":"123456789"}''')
        while True:
            WS.send("ping")
            sleep(10)
    
    _thread.start_new_thread(Run, ())

if __name__ == "__main__":
    strRoomId = "20028484"
    intUnixTs = int(time())
    
    strHuyaToken = GetHuyaToken(strRoomId, 
                                ConstMyTokens.AppId, 
                                ConstMyTokens.AppSecret)

    dicRequestHeader = {"roomId": strRoomId, 
                        "appId": ConstMyTokens.AppId, 
                        "iat": intUnixTs, 
                        "sToken": strHuyaToken}
    
    strRequestURL = "ws://ws-apiext.huya.com/index.html?do=comm&roomId={roomId}&appId={appId}&iat={iat}&sToken={sToken}"
    websocket.enableTrace(True)
    wsHuya = websocket.WebSocketApp(strRequestURL.format(**dicRequestHeader), 
                                    on_message=OnMessage,
                                    on_error=OnError,
                                    on_close=OnClose)

    wsHuya.on_open = OnOpen
    wsHuya.run_forever()